import cv2
import numpy as np
import math
from scipy.optimize import curve_fit
from scipy.signal import savgol_filter
import datetime as dt

from rune import *
from tracker import *


class Line:
    '''用于描述直线的类'''
    def __init__(self,A=0,B=0,C=0) -> None:
        self.A = A
        self.B = B
        self.C = C
    
    def SolveLine(self,point1:np.ndarray,point2:np.ndarray):
        '''根据两点求直线参数A,B,C'''
        self.A = point1[1] - point2[1]
        self.B = point2[0] - point1[0]
        self.C = - self.A*point1[0] - self.B*point1[1]
    
    def SolveNormalLineParameters(self,point:np.ndarray):
        '''根据法线经过的一点求法线参数A,B,C'''
        A = -self.B
        B = self.A
        C = - A*point[0] - B*point[1]
        return (A,B,C)
    
    def SolvePointByX(self,x:float):
        '''根据x求y'''
        if self.B != 0:
            return (-self.C-self.A*x)/self.B
        else:
            return 0

    def DetermineIfSameSide(self,point1:np.ndarray,point2:np.ndarray):
        '''判断两点是否在直线同侧'''
        value1 = self.A*point1[0] + self.B*point1[1] + self.C
        value2 = self.A*point2[0] + self.B*point2[1] + self.C
        res = False
        if value1*value2 > 0:
            res = True
        elif value1== 0 and value2 == 0:
            res = True
        return res


def BigRuneFunc(t, a, omega, c, d):
    return -a/omega * np.cos(omega * (t + c)) + (2.090 - a) * (t + c) + d

def SmallRuneFunc(t, a, b, c):
    return a*(t + b) + c

def RuneSpeedFunc(t, a, omega, c):
    return a* np.sin(omega * (t + c)) + (2.09-a)

def CalcDeltaRunePos(t0 , t1, a, omega, c):
    p0 = -a/omega * np.cos(omega * (t0 + c)) + (2.090-a) * t0
    p1 = -a/omega * np.cos(omega * (t1 + c)) + (2.090-a) * t1
    return p1-p0

class RuneDetector():
    def __init__(self,thresh=100):
        self.binary_thresh = thresh
        # self.binary_thresh = 210
        self.Reset()


    def Reset(self):
        self.result_img_ = None
        self.binary_img_ = None
        self.small_lights_ = []
        self.big_lights_ = []
        self.huge_lights_ = []
        self.circular_lights_ = []
        self.runes_ = []
        
        self.center_ = np.array([0,0]) # 能量机关旋转中心
        self.target_center_ = np.array([0,0]) # 待打击的位置
        self.target_angles_ = [] # （收集的数据）待打击扇叶的角位置
        self.key_angles_ = [] # （收集的数据）连续的角位置（用于拟合位置函数）
        self.predict_angles_ = [] # （预测的数据）待打击扇叶的角位置
        self.predict_diffs_ = [] # 预测的角位置与实际角位置的差值
        self.rune_speeds_ = []
        self.rune_speeds_smoothed_ = []
        self.time_stamps_ = []
        self.start_time_ = dt.datetime.now()
        self.tracker = Tracker()
        self.state = DectorState.NO_ACTION
        self.lose_target_count = 0 # 丢失目标的帧数
        self.rune_params = [1.0, 1.9, 1.0, 1.0]
        self.rune_speed_params = [1.0, 1.9, 1.0]


    def CalcTargetAngle(self,pt:np.ndarray):
        '''计算待打击扇叶的角位置'''
        dy = pt[1] - self.center_[1]
        dx = pt[0] - self.center_[0]
        angle = math.atan2(dy,dx)
        if angle<0:
            angle += math.pi*2
        return angle


    def FitTraceCurve(self):
        '''拟合轨迹曲线'''
        # TODO：对数据进行低通滤波，获取位置变化函数的数据
        
        # 根据数据拟合曲线
        # TODO：分别拟合大符和小符的曲线，根据效果判断使用哪一个
        t = np.array(self.time_stamps_)
        y = np.array(self.rune_speeds_smoothed_)
        p0 = [0.9, 1.9, 0.0]# 初始参数
        try:
            popt, pcov = curve_fit(RuneSpeedFunc, t, y, p0)# 拟合数据
            # popt, pcov = curve_fit(BigRuneFunc, t, y, p0)# 拟合数据
            print('拟合参数：',popt)
            self.rune_speed_params = popt
        except:
            print('拟合失败')

        return


    def CollectData(self):
        if len(self.target_angles_)>=50:
            self.state = DectorState.PREDICTING
            self.FitTraceCurve()
            # return
        
        rune:Rune
        target_center = np.array([0,0])
        for rune in self.runes_:
            if rune.type == RuneType.DISACTIVED:
                target_center = self.GetTargetCenter(rune.rect)
                self.target_center_ = target_center
                break
        target_angle = self.CalcTargetAngle(target_center)
        self.target_angles_.append(target_angle)
        # key_angle = target_angle
        # if len(self.target_angles_)>=1:
        #     d_key_angle = key_angle - self.key_angles_[-1]
        #     if d_key_angle > 30/180*math.pi:
        #         key_angle = round(d_key_angle/(72/180*math.pi))
        if len(self.target_angles_)>=2:
            rune_speed = 50*(self.target_angles_[-1]-self.target_angles_[-2])
            if abs(rune_speed)>2:
                rune_speed = self.rune_speeds_[-1]
            self.rune_speeds_.append(rune_speed)
            self.rune_speeds_smoothed_.append(rune_speed)
            self.time_stamps_.append((dt.datetime.now()-self.start_time_).total_seconds())
            if len(self.time_stamps_)>=2:
                print('f(Hz):',1/(self.time_stamps_[-1]-self.time_stamps_[-2]))
        # for _ in range(10):
        #     self.Filter()
        self.Filter()
        return


    def Filter(self):
        if len(self.rune_speeds_smoothed_)<5 or len(self.rune_speeds_)<5:
            return

        # 滑动平均滤波
        # y = np.array(self.rune_speeds_)
        # window = np.ones(int(5))/float(5)
        # y_smooth = np.convolve(y, window, 'same')
        # self.rune_speeds_smoothed_ = list(y_smooth)
        # SG滤波
        # y = np.array(self.rune_speeds_)
        y = np.array(self.rune_speeds_smoothed_)
        y_smooth = savgol_filter(y, 5, 3)
        self.rune_speeds_smoothed_ = list(y_smooth)
        return


    def Predict(self,dt:float):
        self.CollectData()
        
        a = self.rune_speed_params[0]
        omega = self.rune_speed_params[1]
        b = 2.09 - a
        c = self.rune_speed_params[2]
        t0 = self.time_stamps_[-1]
        t1 = t0 + dt
        dp =CalcDeltaRunePos(t0,t1,a,omega,c)
        
        # 计算旋转半径
        rotation_radius = np.linalg.norm(self.target_center_ - self.center_)
        if rotation_radius < 0.1:
            return
        # 计算预测的打击角度
        predict_pos_angle = self.target_angles_[-1] + dp
        if predict_pos_angle>math.pi*2:
            predict_pos_angle-=math.pi*2
        self.predict_angles_.append(predict_pos_angle)
        
        # if len(self.predict_angles_)>=25:
        #     p0 = self.predict_angles_[-23]
        #     p1 = self.target_angles_[-1]
        #     self.predict_diffs_.append(p1-p0)
            
        # 绘制预测的打击点位置
        x = self.center_[0] + rotation_radius*np.cos(predict_pos_angle)
        y = self.center_[1] + rotation_radius*np.sin(predict_pos_angle)
        cv2.circle(self.result_img_, (int(x),int(y)), 5, (0, 0, 255), -1)
        cv2.putText(self.result_img_,'predict pos',
                    (int(x),int(y)-20),
                    cv2.FONT_HERSHEY_SIMPLEX,0.5,
                    (255,255,255),2)
        return


    def PreprocessImage(self,rgb_img:np.ndarray):
        '''图像预处理'''
        self.result_img_ = rgb_img.copy()#用于显示结果的图像
        #灰度化
        gray_img = cv2.cvtColor(rgb_img,cv2.COLOR_BGR2GRAY)
        
        # #分离通道
        # b,g,r = cv2.split(rgb_img.astype(np.float32))
        # gray_img = r-0.2*g-0.8*b
        # gray_img = np.clip(gray_img, 0, 255)
        # gray_img = gray_img.astype(np.uint8)
        
        #二值化
        ret, binary_img = cv2.threshold(gray_img,self.binary_thresh,255,cv2.THRESH_BINARY)
        #闭运算
        kernel = np.ones((3, 3), np.uint8)
        self.binary_img_ = cv2.morphologyEx(binary_img, cv2.MORPH_CLOSE, kernel,iterations=2)


    def FindLights(self):
        '''寻找图像中的灯条并进行初步的形状分类'''
        # 寻找轮廓
        contours, hierarchy = cv2.findContours(self.binary_img_,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
        # 初始化各个灯条列表
        self.small_lights_.clear()
        self.big_lights_.clear()
        self.huge_lights_.clear()
        self.circular_lights_.clear()
        # 遍历轮廓
        for contour in contours:
            # 排除面积太小的轮廓
            if cv2.contourArea(contour) < 200:
                continue
            # 拟合矩形
            rect = cv2.minAreaRect(contour)
            # 排除长宽比过大的轮廓
            height = max(rect[1][0],rect[1][1])
            width = min(rect[1][0],rect[1][1])
            if height/width > 3:
                continue
            
            # 判断灯条类型
            aspect_ratio = height/width # 计算长宽比
            light = Light(contour,rect)
            if aspect_ratio < 1.1 :# 圆形灯条
                light.type = LightType.CIRCULAR_LIGHT
                self.circular_lights_.append(light)
            elif abs(aspect_ratio-1.25) < 0.1 :# 大型灯条
                light.type = LightType.BIG_LIGHT
                self.big_lights_.append(light)
            elif abs(aspect_ratio-2.7) < 0.2 :# 小型灯条
                light.type = LightType.SMALL_LIGHT
                self.small_lights_.append(light)
            elif abs(aspect_ratio-1.85) < 0.25 :# 巨型灯条
                area = rect[1][0]*rect[1][1]# 计算外接矩形面积
                if area < 20000:# 排除面积过小的矩形
                    continue
                light.type = LightType.HUGE_LIGHT
                self.huge_lights_.append(light)
        # self.DrawLights()


    def DistinguishingLights(self):
        '''对灯条进行进一步的种类分类'''
        self.runes_.clear()
        self.lose_target_count += 1
        # 对分类简单的巨型灯条进行分类
        light:Light
        disactived_huge_light = None
        disactived_is_found = False
        for light in self.huge_lights_:
            # 计算轮廓面积与外接矩形面积的比值
            area_ratio = round(cv2.contourArea(light.contour)/(light.rect[1][0]*light.rect[1][1]),3)
            rune = Rune(light.rect)
            if area_ratio > 0.6:
                rune.type = RuneType.ACTIVED
            elif not disactived_is_found:
                rune.type = RuneType.DISACTIVED
                disactived_huge_light = light
                disactived_is_found = True
                self.lose_target_count = 0
            self.runes_.append(rune)

            

        # 对圆形灯条进行分类（以待激活扇叶坐标为判断依据）
        if disactived_huge_light is None:
            return
        rect = disactived_huge_light.rect
        # 获取待激活扇叶外接矩形四个顶点的坐标
        box = cv2.boxPoints(rect)
        box = np.int0(box)
        if np.linalg.norm(box[0]-box[1]) < np.linalg.norm(box[1]-box[2]):
            line1_points = (box[0],box[1])
            line2_points = (box[2],box[3])
        else:
            line1_points = (box[1],box[2])
            line2_points = (box[0],box[3])
        line1 = Line()
        line1.SolveLine(line1_points[0],line1_points[1])# 求解短边所在的直线

        for light in self.circular_lights_:
            center = light.rect[0]
            normal_line_params = line1.SolveNormalLineParameters(center)# 求解经过center的法线参数
            line1_normal_line = Line(normal_line_params[0],
                                     normal_line_params[1],
                                     normal_line_params[2])# 创建法线
            line2_point_in_same_side = line1_normal_line.DetermineIfSameSide(line2_points[0],line2_points[1])# 判断两条直线是否在同一侧
            
            rune = Rune(light.rect)
            if not line2_point_in_same_side:
                rune.type = RuneType.R_SIGN
                self.runes_.append(rune)
                self.center_ = np.array(center)
                if self.state == DectorState.NO_ACTION:
                    self.state = DectorState.COLLECTING
            else:
                rune.type = RuneType.TargetRing
                self.runes_.append(rune)
            # 在图中绘制line1_normal_line和line1
            # cv2.line(self.result_img_,
            #          (int(center[0]-300),int(line1_normal_line.SolvePointByX(center[0]-300))),
            #          (int(center[0]+300),int(line1_normal_line.SolvePointByX(center[0]+300))),
            #          (255,0,0),2)# 过center的法线
            # cv2.line(self.result_img_,
            #          (int(line1_points[0][0]),int(line1_points[0][1])),
            #          (int(line1_points[1][0]),int(line1_points[1][1])),
            #          (0,255,0),2)# line1
            # cv2.line(self.result_img_,
            #          (int(line2_points[0][0]),int(line2_points[0][1])),
            #          (int(line2_points[1][0]),int(line2_points[1][1])),
            #          (0,255,0),2)# line2
        self.DrawRunes()
        return


    def GetTargetCenter(self,rect:cv2.RotatedRect):
        '''获取打击中心坐标'''
        box = cv2.boxPoints(rect)# 获取旋转矩形的四个顶点
        if np.linalg.norm(box[0]-box[1]) < np.linalg.norm(box[1]-box[2]):
            dst_pts = np.float32([[0, 0], [0, 75], [150, 75]])# 定义仿射变换的目标点
        else:
            dst_pts = np.float32([[0, 75], [150, 75], [150, 0]])# 定义仿射变换的目标点
        
        M = cv2.getAffineTransform(box[:3], dst_pts)# 获取仿射变换矩阵
        M_inv = np.linalg.inv(np.vstack([M, [0, 0, 1]]))# 计算 M 的逆矩阵

        transformed = cv2.warpAffine(self.binary_img_, M, 
                                     (self.binary_img_.shape[1], self.binary_img_.shape[0])
                                    ) # 使用 cv2.warpAffine() 进行仿射变换
        transformed = transformed[0:75,0:150]# 裁剪
        # cv2.imshow('Transformed', transformed)
        extend = 20
        bordered = cv2.copyMakeBorder(transformed, 
                                      extend, extend, extend, extend, 
                                      cv2.BORDER_CONSTANT, value=[0, 0, 0])
        # cv2.imshow('bordered', bordered)
        #闭运算
        kernel = np.ones((15, 15), np.uint8)
        closed = cv2.morphologyEx(bordered, cv2.MORPH_CLOSE, kernel,iterations=2)
        # cv2.imshow('closed', closed)
        
        #开运算
        opened = cv2.morphologyEx(closed, cv2.MORPH_OPEN, kernel,iterations=2)
        # cv2.imshow('opened', opened)

        # 开始寻找打击中心
        contours, hierarchy = cv2.findContours(opened,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
        for contour in contours:
            rect = cv2.minAreaRect(contour)
            
            pt = np.array([
                    [rect[0][0]-extend,rect[0][1]-extend]
                ])#逆变换点pt
            # 对 pt 进行逆仿射变换
            pt_transformed = cv2.perspectiveTransform(np.array([pt], dtype='float32'), M_inv)
            # 绘制pt在原图中的位置
            pos = tuple(pt_transformed[0][0].astype(int))
            cv2.circle(self.result_img_, pos, 5, (0, 255, 0), -1)
            cv2.putText(self.result_img_,'target pos',
                    (pos[0],pos[1]-20),
                    cv2.FONT_HERSHEY_SIMPLEX,0.5,
                    (255,255,255),2)
            return pt_transformed[0][0].astype(int)
        return np.array([0,0])

    def DrawLights(self):
        light:Light
        dy = -15
        
        for light in self.huge_lights_:
            # 绘制外接矩形
            box = cv2.boxPoints(light.rect)
            box = np.int0(box)
            cv2.drawContours(self.result_img_,[box],0,(255,255,0),2)
            # cv2.putText(self.result_img_,'huge light',
            #         (int(light.rect[0][0]),int(light.rect[0][1])+dy),
            #         cv2.FONT_HERSHEY_SIMPLEX,0.5,
            #         (255,255,255),2)
            area = round(light.rect[1][0]*light.rect[1][1],3)# 计算外接矩形面积
            cv2.putText(self.result_img_,'area='+str(area),
                    (int(light.rect[0][0]),int(light.rect[0][1])+dy),
                    cv2.FONT_HERSHEY_SIMPLEX,0.5,
                    (255,255,255),2)
            area_ratio = round(cv2.contourArea(light.contour)/(light.rect[1][0]*light.rect[1][1]),3)
            cv2.putText(self.result_img_,'area_raito='+str(area_ratio),
                    (int(light.rect[0][0]),int(light.rect[0][1])+dy*2),
                    cv2.FONT_HERSHEY_SIMPLEX,0.5,
                    (255,255,255),2)
        return
    
        for light in self.circular_lights_:
            # 绘制外接椭圆
            cv2.ellipse(self.result_img_,light.rect,(0,255,0),2)
            cv2.putText(self.result_img_,'circular light',
                    (int(light.rect[0][0]),int(light.rect[0][1])+dy),
                    cv2.FONT_HERSHEY_SIMPLEX,0.5,
                    (255,255,255),2)
        
        for light in self.big_lights_:
            # 绘制外接矩形
            box = cv2.boxPoints(light.rect)
            box = np.int0(box)
            cv2.drawContours(self.result_img_,[box],0,(255,0,0),2)
            cv2.putText(self.result_img_,'big light',
                    (int(light.rect[0][0]),int(light.rect[0][1])+dy),
                    cv2.FONT_HERSHEY_SIMPLEX,0.5,
                    (255,255,255),2)
        
        for light in self.small_lights_:
            # 绘制外接矩形
            box = cv2.boxPoints(light.rect)
            box = np.int0(box)
            cv2.drawContours(self.result_img_,[box],0,(0,0,255),2)
            cv2.putText(self.result_img_,'small light',
                    (int(light.rect[0][0]),int(light.rect[0][1])+dy),
                    cv2.FONT_HERSHEY_SIMPLEX,0.5,
                    (255,255,255),2)
        


    def DrawRunes(self):
        rune:Rune
        dy = 10
        for rune in self.runes_:
            box = cv2.boxPoints(rune.rect)
            box = np.int0(box)
            cv2.drawContours(self.result_img_,[box],0,(255,255,0),2)
            
            if rune.type == RuneType.DISACTIVED:
                cv2.putText(self.result_img_,'disactived',
                        (int(rune.rect[0][0]),int(rune.rect[0][1])+dy),
                        cv2.FONT_HERSHEY_SIMPLEX,0.5,
                        (255,255,255),2)
            elif rune.type == RuneType.ACTIVED:
                cv2.putText(self.result_img_,'actived',
                        (int(rune.rect[0][0]),int(rune.rect[0][1])+dy),
                        cv2.FONT_HERSHEY_SIMPLEX,0.5,
                        (255,255,255),2)
            elif rune.type == RuneType.R_SIGN:
                cv2.putText(self.result_img_,'R sign',
                        (int(rune.rect[0][0]),int(rune.rect[0][1])-20),
                        cv2.FONT_HERSHEY_SIMPLEX,0.5,
                        (0,255,255),2)
            elif rune.type == RuneType.TargetRing:
                cv2.putText(self.result_img_,'target ring',
                        (int(rune.rect[0][0]),int(rune.rect[0][1])+dy),
                        cv2.FONT_HERSHEY_SIMPLEX,0.5,
                        (0,255,255),2)
            else:
                cv2.putText(self.result_img_,'invalid',
                        (int(rune.rect[0][0]),int(rune.rect[0][1])+dy),
                        cv2.FONT_HERSHEY_SIMPLEX,0.5,
                        (255,255,255),2)



        
        
































